# -*- coding: utf-8 -*-
"""
Redis 分层时间轮实现
===================

这是一个基于 Redis 的分布式时间轮实现，用于高效管理和调度定时任务。

主要特性:
- 多层时间轮结构，支持不同时间精度的调度
- 基于 Redis 实现，支持分布式部署
- 默认支持最长24小时的调度周期，超出部分存入溢出桶
- 固定使用秒级精度的任务调度
- 使用分布式锁确保并发安全

注意事项:
1. 时间处理
   - 所有时间默认使用系统本地时区
   - 如果传入无时区信息的datetime，会被视为本地时间
   - 内部存储使用秒级时间戳
   -r

2. Redis键前缀
   - 所有Redis键都使用"time_wheel:"前缀
   - 请确保Redis实例中该前缀未被其他服务使用

3. 性能考虑
   - 合理设置wheel_size和layers参数
   - 默认配置适用于大多数场景，支持:
     * 第1层: 1分钟 (1秒 * 60)
     * 第2层: 1小时
     * 第3层: 1天
     * 第4层: 15天

4. 分布式部署
   - 使用DistributedLock确保并发安全
   - 多实例部署时需确保Redis连接配置一致
   - 建议使用Redis集群保证高可用

维护说明:
- 定期检查Redis内存使用情况
- 监控time_wheel:overflow中的任务数量
- 建议设置Redis键的过期时间
- 定期清理过期的锁键
"""
from datetime import datetime, timezone
from typing import List, Tuple

from app.core.config import settings
from app.core.redis_conn import redis_helper
from app.utils.distributed_lock import DistributedLock
from app.utils.timezone import convert_to_local, convert_to_utc, get_timezone, now


class TimeWheelLayer:
    """时间轮层级"""

    def __init__(self, tick_ms: int, wheel_size: int, layer_index: int = 0):
        """
        初始化时间轮层级

        Args:
            tick_ms: 每个槽位代表的时间（毫秒）
            wheel_size: 时间轮大小（槽位数量）
            layer_index: 层级索引，0表示最底层
        """
        self.tick_ms = tick_ms
        self.wheel_size = wheel_size
        self.interval_ms = tick_ms * wheel_size  # 整个时间轮表示的时间范围
        self.layer_index = layer_index

    def get_slot(self, timestamp_ms: int) -> int:
        """获取时间戳对应的槽位"""
        # 使用向下取整确保时间在槽位范围内
        return max(0, (timestamp_ms // self.tick_ms) % self.wheel_size)

    def get_next_slot(self, timestamp_ms: int) -> int:
        """获取下一个槽位"""
        current_slot = self.get_slot(timestamp_ms)
        return (current_slot + 1) % self.wheel_size


class RedisTimeWheel:
    """
    基于Redis的分层时间轮

    实现了一个分布式的多层时间轮，用于高效管理和调度定时任务。
    - 支持多个层级的时间轮，每层精度不同
    - 使用Redis作为存储后端，支持分布式部署
    - 支持添加、移除、查询和执行到期任务
    """

    def __init__(
        self,
        wheel_size: int = 60,
        layers: int = 4,
    ):
        """初始化时间轮

        Args:
            redis_client: Redis客户端实例
            wheel_size: 时间轮大小（槽位数量），默认60
            layers: 时间轮层数，默认4层
        """
        self.redis_client = redis_helper.get_connection(settings.TASK_REDIS_DB)
        self.wheel_size = wheel_size
        self.layers = []
        self.base_tick_ms = 1000  # 固定为1秒的基础时间粒度

        # 初始化各层时间轮
        for i in range(layers):
            tick = self.base_tick_ms * (wheel_size**i)
            layer = TimeWheelLayer(tick, wheel_size, i)
            self.layers.append(layer)

        # 计算最大时间范围（最高层间隔）
        self.total_range_ms = self.layers[-1].interval_ms if self.layers else 0

        # Redis键前缀
        self.key_prefix = "time_wheel:"
        self.current_time_key = f"{self.key_prefix}current_time"

    def _get_layer_key(self, layer_index: int) -> str:
        """获取层级对应的Redis键"""
        return f"{self.key_prefix}layer:{layer_index}"

    def _get_overflow_key(self) -> str:
        """获取溢出桶的Redis键"""
        return f"{self.key_prefix}overflow"

    def _timestamp_ms(self, dt: datetime) -> int:
        """将datetime转换为毫秒时间戳"""
        # 确保时间有时区信息
        if dt.tzinfo is None:
            dt = get_timezone().localize(dt)
        # 转换为UTC时间戳
        utc_dt = convert_to_utc(dt)
        return int(utc_dt.timestamp() * 1000)

    def _from_timestamp_ms(self, ts_ms: int) -> datetime:
        """将毫秒时间戳转换为datetime"""
        # 使用 UTC 时间戳创建 datetime 对象
        utc_dt = datetime.fromtimestamp(ts_ms / 1000, timezone.utc)
        # 转换为本地时区
        local_dt = convert_to_local(utc_dt)
        # 保持毫秒精度
        microseconds = int((ts_ms % 1000) * 1000)
        return local_dt.replace(microsecond=microseconds)

    async def add_task(self, task_id: str, execute_time: datetime) -> None:
        """添加任务到时间轮"""
        if execute_time.tzinfo is None:
            execute_time = get_timezone().localize(execute_time)

        current_time = await self.get_current_time()
        if current_time.tzinfo is None:
            current_time = get_timezone().localize(current_time)

        execute_ms = self._timestamp_ms(execute_time)
        current_ms = self._timestamp_ms(current_time)
        time_diff_ms = execute_ms - current_ms

        # 如果执行时间超过24小时，放入溢出桶
        ONE_DAY_MS = 24 * 60 * 60 * 1000
        if time_diff_ms >= ONE_DAY_MS:
            async with DistributedLock(self.redis_client, f"lock:add_task:{task_id}"):
                await self.redis_client.zadd(self._get_overflow_key(), {task_id: execute_ms})
            return

        # 找到合适的层级
        target_layer = None
        for i, layer in enumerate(self.layers):
            if time_diff_ms < layer.interval_ms:
                target_layer = layer
                break

        async with DistributedLock(self.redis_client, f"lock:add_task:{task_id}"):
            if target_layer:
                slot = target_layer.get_slot(execute_ms)
                key = self._get_layer_key(target_layer.layer_index)
                # 使用 Redis 的 HSETNX 命令确保不会覆盖已存在的任务
                await self.redis_client.hsetnx(f"{key}:{slot}", task_id, str(execute_ms))
            else:
                await self.redis_client.zadd(self._get_overflow_key(), {task_id: execute_ms})

    async def get_due_tasks(self) -> List[Tuple[str, datetime]]:
        """获取并移除所有到期的任务"""
        current_time = await self.get_current_time()
        if current_time.tzinfo is None:
            current_time = get_timezone().localize(current_time)

        current_ms = self._timestamp_ms(current_time)
        due_tasks = []

        async with DistributedLock(self.redis_client, "lock:get_due_tasks"):
            # 检查所有层级
            for layer_index, layer in enumerate(self.layers):

                # 检查所有槽位中的任务
                for slot in range(layer.wheel_size):
                    slot_key = f"{self._get_layer_key(layer_index)}:{slot}"
                    tasks = await self.redis_client.hgetall(slot_key)

                    if tasks:

                        # 批量处理到期任务
                        expired_tasks = []
                        for task_id_raw, execute_ms_raw in tasks.items():
                            task_id = task_id_raw.decode() if isinstance(task_id_raw, bytes) else str(task_id_raw)
                            execute_ms = int(
                                execute_ms_raw.decode() if isinstance(execute_ms_raw, bytes) else str(execute_ms_raw)
                            )
                            execute_time = self._from_timestamp_ms(execute_ms)

                            if execute_ms <= current_ms:
                                expired_tasks.append((task_id_raw, execute_time))

                        # 批量删除到期任务
                        if expired_tasks:
                            await self.redis_client.hdel(slot_key, *[task[0] for task in expired_tasks])
                            due_tasks.extend(
                                [
                                    (task_id.decode() if isinstance(task_id, bytes) else str(task_id), time)
                                    for task_id, time in expired_tasks
                                ]
                            )

            # 检查溢出桶
            overflow_tasks = await self.redis_client.zrangebyscore(
                self._get_overflow_key(), 0, current_ms, withscores=True
            )

            if overflow_tasks:
                await self.redis_client.zremrangebyscore(self._get_overflow_key(), 0, current_ms)
                for task_id_raw, score in overflow_tasks:
                    task_id = task_id_raw.decode() if isinstance(task_id_raw, bytes) else str(task_id_raw)
                    execute_time = self._from_timestamp_ms(int(score))
                    due_tasks.append((task_id, execute_time))

        return due_tasks

    async def advance(self) -> None:
        """推进时间轮的当前时间"""
        new_time = now()
        if new_time.tzinfo is None:
            new_time = get_timezone().localize(new_time)

        current_ms = self._timestamp_ms(new_time)

        async with DistributedLock(self.redis_client, "lock:advance_time"):
            # 获取旧的时间
            old_time_str = await self.redis_client.get(self.current_time_key)
            if old_time_str:
                try:
                    old_time = datetime.fromisoformat(old_time_str.decode())
                    if old_time.tzinfo is None:
                        old_time = get_timezone().localize(old_time)
                    old_ms = self._timestamp_ms(old_time)
                except (ValueError, AttributeError):
                    old_ms = current_ms
            else:
                old_ms = current_ms

            # 如果时间差超过一个槽位，需要级联重新分配任务
            time_diff_ms = current_ms - old_ms
            if time_diff_ms >= self.layers[0].tick_ms:
                # 检查每一层是否需要重新分配任务
                for i, layer in enumerate(self.layers):
                    if time_diff_ms >= layer.tick_ms:
                        old_slot = layer.get_slot(old_ms)
                        new_slot = layer.get_slot(current_ms)

                        # 如果槽位发生变化，需要重新分配任务
                        if old_slot != new_slot:
                            # 获取旧槽位的所有任务
                            slot_key = f"{self._get_layer_key(i)}:{old_slot}"
                            tasks = await self.redis_client.hgetall(slot_key)

                            # 批量处理任务
                            to_delete = []
                            to_move = {}
                            to_readd = []

                            for task_id_raw, execute_ms_raw in tasks.items():
                                task_id = task_id_raw.decode() if isinstance(task_id_raw, bytes) else str(task_id_raw)
                                execute_ms = int(
                                    execute_ms_raw.decode()
                                    if isinstance(execute_ms_raw, bytes)
                                    else str(execute_ms_raw)
                                )

                                # 如果任务已经到期，加入删除列表
                                if execute_ms <= current_ms:
                                    to_delete.append(task_id_raw)
                                    # 将到期任务添加到临时存储中
                                    await self.redis_client.hset(
                                        f"{self._get_layer_key(0)}:expired", task_id, str(execute_ms)
                                    )
                                    continue

                                # 重新计算时间差
                                time_diff = execute_ms - current_ms

                                # 如果时间差小于当前层级的间隔，加入移动列表
                                if time_diff < layer.interval_ms:
                                    new_slot_key = f"{self._get_layer_key(i)}:{new_slot}"
                                    to_move[task_id] = str(execute_ms)
                                else:
                                    # 否则需要重新添加到合适的层级
                                    to_readd.append((task_id, self._from_timestamp_ms(execute_ms)))

                            # 批量执行操作
                            if to_delete:
                                await self.redis_client.hdel(slot_key, *to_delete)
                            if to_move:
                                new_slot_key = f"{self._get_layer_key(i)}:{new_slot}"
                                await self.redis_client.hmset(new_slot_key, to_move)
                            for task_id, execute_time in to_readd:
                                await self.add_task(task_id, execute_time)

                            # 清理旧槽位
                            await self.redis_client.delete(slot_key)

            # 更新当前时间
            await self.redis_client.set(self.current_time_key, new_time.isoformat())

    async def _cascade_tasks(self, old_ms: int, new_ms: int) -> None:
        """
        级联重新分配任务

        当时间推进超过一个槽位时，需要将高层级的任务重新分配到低层级
        """
        for i in range(1, len(self.layers)):
            layer = self.layers[i]
            old_slot = layer.get_slot(old_ms)
            new_slot = layer.get_slot(new_ms)

            if old_slot != new_slot:
                # 获取旧槽位的所有任务
                slot_key = f"{self._get_layer_key(i)}:{old_slot}"
                tasks = await self.redis_client.hgetall(slot_key)

                # 重新分配任务到低层级
                for task_id_raw, execute_ms_raw in tasks.items():
                    # 安全地处理任务ID
                    task_id = task_id_raw.decode() if isinstance(task_id_raw, bytes) else str(task_id_raw)

                    # 安全地处理执行时间
                    if isinstance(execute_ms_raw, bytes):
                        execute_ms = int(execute_ms_raw.decode())
                    else:
                        execute_ms = int(str(execute_ms_raw))

                    execute_time = self._from_timestamp_ms(execute_ms)

                    # 如果任务已经到期，直接添加到结果中
                    current_time = await self.get_current_time()
                    if execute_time <= current_time:
                        continue

                    # 否则重新分配到合适的层级
                    await self.add_task(task_id, execute_time)

                # 清空旧槽位
                if tasks:
                    await self.redis_client.delete(slot_key)

    async def get_current_time(self) -> datetime:
        """获取时间轮的当前时间"""
        time_str = await self.redis_client.get(self.current_time_key)
        if time_str:
            try:
                dt = datetime.fromisoformat(time_str.decode())
                return convert_to_local(dt)
            except (ValueError, AttributeError):
                pass
        return now()

    async def count_tasks(self) -> int:
        """
        获取时间轮中的任务总数

        Returns:
            任务总数
        """
        total = 0

        # 统计所有层级的任务数
        for i in range(len(self.layers)):
            for slot in range(self.layers[i].wheel_size):
                slot_key = f"{self._get_layer_key(i)}:{slot}"
                total += await self.redis_client.hlen(slot_key)

        # 统计溢出桶中的任务数
        total += await self.redis_client.zcard(self._get_overflow_key())

        return total

    async def remove_task(self, task_id: str) -> bool:
        """
        从时间轮中移除任务

        Args:
            task_id: 任务ID

        Returns:
            是否成功移除任务
        """
        removed = False

        async with DistributedLock(self.redis_client, f"lock:remove_task:{task_id}"):
            # 检查所有层级
            for i in range(len(self.layers)):
                for slot in range(self.layers[i].wheel_size):
                    slot_key = f"{self._get_layer_key(i)}:{slot}"
                    if await self.redis_client.hexists(slot_key, task_id):
                        await self.redis_client.hdel(slot_key, task_id)
                        removed = True
                        break
                if removed:
                    break

            # 如果在层级中没有找到，检查溢出桶
            if not removed:
                removed = await self.redis_client.zrem(self._get_overflow_key(), task_id) > 0

        return removed

    async def clear_tasks(self) -> None:
        """清空所有任务

        使用分布式锁保证并发安全，并通过 pipeline 批量删除提升性能
        """
        async with DistributedLock(self.redis_client, "lock:clear_tasks"):
            # 使用 pipeline 批量执行删除操作
            pipe = self.redis_client.pipeline()

            # 删除当前时间和溢出桶
            pipe.delete(self.current_time_key)
            pipe.delete(self._get_overflow_key())

            # 批量删除所有层级的槽位
            for layer in self.layers:
                keys = [f"{self._get_layer_key(layer.layer_index)}:{slot}" for slot in range(layer.wheel_size)]
                if keys:
                    pipe.delete(*keys)

            # 执行所有删除操作
            await pipe.execute()
